#!/usr/bin/python3

import os
import subprocess
import sys
import signal
import argparse
from pathlib import Path
from subprocess import CompletedProcess
from typing import Dict, List, Tuple

# 全局变量，用于跟踪测试进度和失败用例
tests_run = 0
tests_passed = 0
failed_tests: List[Tuple[str, str]] = []

REASON_MAPPING = {
    "COMPILATION_ERROR": {"color": "33", "desc": "Compilation Error"},  # 黄色
    "OUTPUT_MISMATCH": {"color": "31", "desc": "Output Mismatch"},  # 红色
    "MISSING_EXPECTED": {"color": "35", "desc": "Missing Expected"}  # 紫色
}

compiler_path: str = os.environ["IRIS_CC"]
llvm_link_path: str = os.environ["IRIS_LLVM_LINK"]
sysy_lib_path: str = os.environ['IRIS_SYSY_LIB']
modified_llc_path: str = os.environ['IRIS_LLC']
riscv32_sysroot: str = os.environ['RV32_SYSROOT']
riscv32_cross_compiler: str = os.environ['RV32_GCC']
qemu_riscv32: str = os.environ['QEMU_RV32']
sysy_h : str = os.environ['IRIS_SYSY_HEADER']
sysy_c : str = os.environ['IRIS_SYSY_C']


def signal_handler(sig, frame):
    """处理 Ctrl+C 信号，打印当前测试进度并退出"""
    global failed_tests
    print(f"\n\033[33m[INTERRUPTED] Testing interrupted by Ctrl+C\033[0m")
    print(f"Tests run so far: {tests_run}")
    print(f"Tests passed so far: {tests_passed}")
    print(f"Pass rate so far: {tests_passed / tests_run * 100:.2f}%" if tests_run > 0 else "Pass rate so far: N/A")

    if failed_tests:
        print("\nFailed tests details:")
        for path, reason in failed_tests:
            info = REASON_MAPPING.get(reason, {"color": "37", "desc": reason})
            abs_path = str(Path(path).absolute())  # 确保绝对路径
            print(f"  \033[{info['color']}m{abs_path}\033[0m - {info['desc']}")
    else:
        print("\nNo failed tests so far.")
    sys.exit(0)


def run_test(test_file: Path) -> str:
    global tests_run, tests_passed
    test_dir: Path = test_file.parent
    test_basename: str = test_file.stem
    expected_output_file: Path = test_dir / f"{test_basename}.out"
    input_file: Path = test_dir / f"{test_basename}.in"

    if not expected_output_file.exists():
        print(f"\033[33m[WARN] Missing expected output file for {test_file}\033[0m")
        tests_run += 1
        return "MISSING_EXPECTED"

    # Compile C to LLVM IR
    compile_proc: subprocess.CompletedProcess = subprocess.run(
        [compiler_path, test_file], capture_output=True, text=True
    )

    if compile_proc.returncode != 0:
        print(f"\033[33m[FAIL] Compilation failed: {test_file}\033[0m")  # 保持黄色警告
        tests_run += 1
        return "COMPILATION_ERROR"

    llvm_ir_code : str = compile_proc.stdout
    llvm_ir_code = llvm_ir_code.replace("inttoptr (i64 3735928559 to ptr)", "null")

    open("tmp.ll", "wt").write(llvm_ir_code)

    # llvm_link_proc: subprocess.CompletedProcess = subprocess.run(
    #     [llvm_link_path, sysy_lib_path, "tmp.ll", "-S", "-o", "out.ll"],
    #     capture_output=True,
    #     text=True
    # )

    llc_proc: subprocess.CompletedProcess = subprocess.run(
        [modified_llc_path, "--march=riscvn", "tmp.ll", "-o", "out.s"],
        # stdout=subprocess.DEVNULL,
        # stderr=subprocess.DEVNULL,
    )

    if llc_proc.returncode != 0:
        print(f"\033[33m[FAIL] Compilation failed: {test_file}\033[0m")  # 保持黄色警告
        tests_run += 1
        return "LLC_COMPILATION_ERROR"

    cross_compiler_proc: CompletedProcess[bytes] = subprocess.run(
        [riscv32_cross_compiler, "out.s", "-o", "out.exe", sysy_h, sysy_c]
    )

    if cross_compiler_proc.returncode != 0:
        print(f"\033[33m[FAIL] Compilation failed: {test_file}\033[0m")  # 保持黄色警告
        tests_run += 1
        return "CROSS_COMPILER_COMPILATION_ERROR"

    # Run with lli
    input_data: str | None = input_file.read_text() if input_file.exists() else None
    # lli_proc: subprocess.CompletedProcess = subprocess.run(
    #     [lli_path, "out.ll"],
    #     input=input_data,
    #     capture_output=True,
    #     text=True
    # )

    qemu_proc = subprocess.run(
        [qemu_riscv32, "-L", riscv32_sysroot, "out.exe"],
        input=input_data,
        capture_output=True,
        text=True
    )

    # output: str = lli_proc.stdout.strip()
    # return_code: int = lli_proc.returncode

    output: str = qemu_proc.stdout.strip()
    return_code: int = qemu_proc.returncode

    # Read expected output
    with open(expected_output_file, "r") as f:
        expected_lines: list[str] = f.read().strip().split("\n")
    expected_output: str = "\n".join(expected_lines[:-1]).strip()
    expected_return: int = int(expected_lines[-1])

    # Compare results
    tests_run += 1
    if output == expected_output and return_code == expected_return:
        print(f"\033[32m[PASS] {test_file}\033[0m")
        tests_passed += 1
        return "PASS"
    else:
        print(f"\033[31m[FAIL] {test_file}\033[0m")  # 输出不匹配保持红色
        print(f"  Expected Output: {expected_output}")
        print(f"  Actual Output:   {output}")
        print(f"  Expected Return: {expected_return}")
        print(f"  Actual Return:   {return_code}")
        return "OUTPUT_MISMATCH"


def main() -> None:
    global failed_tests

    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="Run tests with optional number filter")
    parser.add_argument("-n", "--number", type=str, help="Run only tests starting with this number")
    args = parser.parse_args()

    # 注册信号处理程序
    signal.signal(signal.SIGINT, signal_handler)

    # compiler_path: str = os.environ["IRIS_CC"]
    # lli_path: str = os.environ["IRIS_LLI"]
    # llvm_link_path: str = os.environ["IRIS_LLVM_LINK"]
    # sysy_lib_path: str = os.environ['IRIS_SYSY_LIB']
    # modified_llc_path: str = os.environ['IRIS_LLC']
    # riscv32_sysroot : str = os.environ['RV32_SYSROOT']
    # riscv32_cross_compiler: str = os.environ['RV32_GCC']
    test_root: Path = Path("test")

    # 获取测试文件列表
    test_files: list[Path] = list(test_root.glob("*.c"))
    test_files.sort()

    # 如果指定了-n参数，过滤测试文件
    if args.number:
        test_files = [f for f in test_files if f.stem.startswith(args.number)]
        if not test_files:
            print(f"\033[33m[WARN] No test files found starting with '{args.number}'\033[0m")
            sys.exit(1)

    total_tests: int = len(test_files)
    results: Dict[str, int] = {"PASS": 0, "COMPILATION_ERROR": 0, "OUTPUT_MISMATCH": 0, "MISSING_EXPECTED": 0,
                               "LLC_COMPILATION_ERROR": 0, "CROSS_COMPILER_COMPILATION_ERROR": 0}

    for test_file in test_files:
        abs_path = test_file.absolute()  # 获取绝对路径
        result: str = run_test(test_file)
        results[result] += 1
        if result != "PASS":
            failed_tests.append((str(abs_path), result))  # 存储绝对路径

    print("\nTest Summary:")
    print(f"  \033[32mPASS: {results['PASS']}\033[0m")
    print(f"  \033[33mCompilation Errors: {results['COMPILATION_ERROR']}\033[0m")
    print(f"  \033[31mOutput Mismatch: {results['OUTPUT_MISMATCH']}\033[0m")
    print(f"  \033[35mMissing Expected Output: {results['MISSING_EXPECTED']}\033[0m")
    print(f"  Total Tests: {total_tests}, Pass Rate: {results['PASS'] / total_tests * 100:.2f}%")

    if failed_tests:
        print("\nFailed tests details:")
        for path, reason in failed_tests:
            info = REASON_MAPPING.get(reason, {"color": "37", "desc": reason})
            print(f"  \033[{info['color']}m{path}\033[0m - {info['desc']}")
    else:
        print("\nAll tests passed!")


if __name__ == "__main__":
    main()
